# 获取日k线
import baostock as bs
import pandas as pd
import time
from datetime import timedelta, datetime


def split_time_interval(start_time, end_time, num_intervals):
    """
    将从start_time到end_time的时间段等分为num_intervals个时间段
    """
    total_duration = end_time - start_time
    interval_duration = total_duration / num_intervals
    
    for i in range(num_intervals):
        interval_start = start_time + i * interval_duration
        interval_end = interval_start + interval_duration
        a = f"{interval_start.year}-{interval_start.month}-{interval_start.day}"
        b = f"{interval_end.year}-{interval_end.month}-{interval_end.day}"
        yield (a, b)


#### 登陆系统 ####
lg = bs.login()
# 显示登陆返回信息
print('login respond error_code:'+lg.error_code)
print('login respond  error_msg:'+lg.error_msg)

#### 获取沪深A股历史K线数据 ####
# 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。“分钟线”不包含指数。
# 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
# 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
params = {
    'id': "sh.600438",    # 通威股份的
    # 'id': "sh.600519", # 茅台的
    'fields': "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",
    'start_date': "2023-06-22",
    'end_date': "2023-12-22",
    'frequency': "d",
    # 复权类型，默认不复权：3；1：后复权；2：前复权。
    'adjustflag': "2",
}

# 设置查询开始和结束日期
start = datetime(2024, 1, 1)  # 2023年3月1日
end = datetime(2024, 7, 18)     # 2023年3月2日
num_intervals = 3            # 设置分为几次查询

list = []

for start_time, end_time in split_time_interval(start, end, num_intervals):
    list.append({ 'start_date': start_time, 'end_date': end_time, })

len = len(list)
for index, value in enumerate(list):
    if(index < len - 1):
        if(list[index + 1]['start_date'] == list[index]['end_date']):
            date = datetime.strptime(list[index + 1]['start_date'], "%Y-%m-%d") + timedelta(days=1)
            list[index + 1]['start_date'] = f"{date.year}-{date.month}-{date.day}"

print(list)
# exit(0)   # 终止程序


data_list = []
columnsDemo = ['date','code','open','high','low','close','preclose','volume','amount','adjustflag','turn','tradestatus','pctChg','isST']

def getData(params):
    rs = bs.query_history_k_data_plus(
        params['id'],
        params['fields'],
        start_date = params['start_date'],
        end_date = params['end_date'],
        frequency = params['frequency'],
        adjustflag = params['adjustflag'],
    )
    print('query_history_k_data_plus respond error_code:'+rs.error_code)
    print('query_history_k_data_plus respond error_msg:'+rs.error_msg)

    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())


for i in list:
    print('------------------------')
    print(i)
    params['start_date'] = i['start_date']
    params['end_date'] = i['end_date']
    getData(params)
    time.sleep(3)

result = pd.DataFrame(data_list, columns= columnsDemo)


#### 结果集输出到csv文件 ####   
result.to_csv("./files/通威股份.csv", index=False)
# result.to_csv("././files/茅台.csv", index=False)
print(result)

#### 登出系统 ####
bs.logout()
